#!/usr/bin/env python3
"""
文件整理助手 - 基于LangChain的智能文件整理系统
通过大语言模型理解自然语言指令，安全整理指定目录下的文件

配置说明:
1. 可以通过命令行交互式输入API配置信息
2. 也可以在同目录下的 .env 文件中配置以下环境变量:
   - API_BASE: API基础URL (例如 https://api.example.com/v1)
   - MODEL_NAME: 模型名称 (例如 qwen-max, gpt-4, mistral-large)
   - API_KEY: API密钥
"""

import os
import sys
import json
import shutil
import datetime
from typing import List, Dict, Any, Optional, Callable
from getpass import getpass

# 添加 dotenv 导入
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    print("提示: 未安装python-dotenv。将使用交互式方式获取API配置。")
    print("建议执行: pip install python-dotenv")

# 尝试导入LangChain相关模块，如果没有安装则提示安装
try:
    from langchain.agents import AgentExecutor, create_tool_calling_agent
    from langchain.tools import Tool, StructuredTool
    from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
    from langchain_core.callbacks import CallbackManagerForToolRun
    from langchain_community.callbacks import get_openai_callback
    from langchain_openai import ChatOpenAI
    from pydantic import BaseModel, Field
except ImportError:
    print("错误: 未安装LangChain相关包。请先安装依赖。")
    print("建议执行: pip install langchain langchain-community langchain-core langchain-openai")
    sys.exit(1)

class FileSystemMCP:
    """文件系统管理控制面板 - 作为大模型可以调用的工具"""
    
    def __init__(self, base_dir: str):
        self.base_dir = os.path.abspath(base_dir)
        self.operation_log = []
        self.planned_operations = []
        print(f"✅ 文件系统MCP已初始化，工作目录: {self.base_dir}")
    
    def list_directory(self, path: str = "") -> str:
        """列出目录内容"""
        full_path = os.path.join(self.base_dir, path)
        if not os.path.exists(full_path):
            return f"❌ 错误: 路径 '{full_path}' 不存在"
        
        if not os.path.isdir(full_path):
            return f"❌ 错误: '{full_path}' 不是目录"
        
        try:
            items = os.listdir(full_path)
            files = []
            dirs = []
            
            for item in items:
                item_path = os.path.join(full_path, item)
                if os.path.isdir(item_path):
                    dirs.append(item)
                else:
                    files.append(item)
            
            result = f"📁 目录内容 ({full_path}):\n"
            if dirs:
                result += "📂 子目录:\n" + "\n".join(f"  • [DIR] {d}" for d in dirs) + "\n"
            if files:
                result += "📄 文件:\n" + "\n".join(f"  • [FILE] {f}" for f in files)
            
            return result
        except Exception as e:
            return f"❌ 列出目录时出错: {str(e)}"
    
    def get_file_structure(self, max_depth: int = 3, path: str = "") -> str:
        """获取目录结构（递归）"""
        full_path = os.path.join(self.base_dir, path)
        if not os.path.exists(full_path) or not os.path.isdir(full_path):
            return f"❌ 错误: 路径 '{full_path}' 不存在或不是目录"
        
        def _build_structure(current_path, current_depth, prefix=""):
            if current_depth > max_depth:
                return []
            
            try:
                items = sorted(os.listdir(current_path))
                structure = []
                
                # 先添加目录
                for item in items:
                    item_path = os.path.join(current_path, item)
                    if os.path.isdir(item_path):
                        structure.append(f"{prefix}📁 {item}")
                        structure.extend(_build_structure(item_path, current_depth + 1, prefix + "  "))
                
                # 再添加文件
                for item in items:
                    item_path = os.path.join(current_path, item)
                    if not os.path.isdir(item_path):
                        # 获取文件大小
                        size = os.path.getsize(item_path)
                        size_str = self._format_size(size)
                        structure.append(f"{prefix}📄 {item} ({size_str})")
                
                return structure
            except Exception as e:
                return [f"❌ 错误: {str(e)}"]
        
        structure = _build_structure(full_path, 0)
        return "\n".join(structure)
    
    def _format_size(self, size: int) -> str:
        """格式化文件大小"""
        for unit in ['B', 'KB', 'MB', 'GB']:
            if size < 1024.0:
                return f"{size:.1f}{unit}"
            size /= 1024.0
        return f"{size:.1f}TB"
    
    def plan_organize(self, organization_plan: str) -> str:
        """制定文件整理计划"""
        try:
            # 检查是否已经是Python对象（列表），如果是则直接使用，否则解析JSON
            if isinstance(organization_plan, list):
                plan = organization_plan
            elif isinstance(organization_plan, str):
                # 如果是字符串，则尝试解析JSON
                plan = json.loads(organization_plan)
            else:
                raise ValueError("无效的整理计划格式")
            
            # 验证计划格式
            if not isinstance(plan, list):
                raise ValueError("整理计划必须是列表格式")
            
            for item in plan:
                if not isinstance(item, dict) or 'source' not in item or 'target' not in item:
                    raise ValueError("每个计划项必须包含'source'和'target'字段")
            
            self.organization_plan = plan
            
            # 生成计划预览
            preview_lines = ["📋 整理计划预览:"]
            for i, item in enumerate(plan, 1):
                preview_lines.append(f"{i}. 将 '{item['source']}' 移动到 '{item['target']}'")
            
            preview_lines.append("\n💡 注意: 此计划尚未执行，需要您的确认")
            return "\n".join(preview_lines)
            
        except json.JSONDecodeError as e:
            return f"❌ JSON解析错误: {str(e)}"
        except Exception as e:
            return f"❌ 制定整理计划时出错: {str(e)}"
    
    def execute_organize(self, confirm: str) -> str:
        """执行文件整理"""
        if not self.organization_plan:
            return "❌ 没有找到整理计划，请先制定整理计划。"
        
        if confirm.lower() not in ['yes', 'y', '是']:
            self.organization_plan = []  # 清除计划
            return "❌ 操作已取消，整理计划已清除。"
        
        results = []
        moved_count = 0
        
        try:
            for item in self.organization_plan:
                source = item['source']
                target = item['target']
                
                # 构建完整路径
                source_path = os.path.join(self.base_dir, source)
                target_path = os.path.join(self.base_dir, target)
                
                # 检查源文件是否存在
                if not os.path.exists(source_path):
                    results.append(f"⚠️ 跳过: 源文件 '{source_path}' 不存在")
                    continue
                
                # 创建目标目录
                target_dir = os.path.dirname(target_path)
                if target_dir and not os.path.exists(target_dir):
                    os.makedirs(target_dir)
                
                # 移动文件
                if os.path.exists(target_path):
                    results.append(f"⚠️ 跳过: 目标文件 '{target}' 已存在")
                else:
                    shutil.move(source_path, target_path)
                    results.append(f"✅ 已移动: '{source}' → '{target}'")
                    moved_count += 1
            
            # 清除计划
            self.organization_plan = []
            
            # 添加总结信息
            results.append(f"\n✅ 整理操作已完成!")
            results.append(f"📊 共移动 {moved_count} 个文件")
            return "\n".join(results)
            
        except Exception as e:
            return f"❌ 执行整理计划时出错: {str(e)}"
    
    def generate_report(self, *args, **kwargs) -> str:
        """生成整理报告"""
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        report_filename = f"整理报告_{timestamp}.txt"
        report_path = os.path.join(self.base_dir, report_filename)
        
        try:
            # 获取当前目录结构
            current_structure = self.get_file_structure(max_depth=3)
            
            with open(report_path, "w", encoding="utf-8") as f:
                f.write(f"===== 文件整理报告 =====\n")
                f.write(f"生成时间: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                f.write(f"工作目录: {self.base_dir}\n")
                f.write("=" * 50 + "\n\n")
                
                f.write("📌 整理后的目录结构:\n")
                f.write(current_structure)
                f.write("\n\n")
                
                f.write("📝 操作日志:\n")
                for entry in self.operation_log:
                    if entry["type"] == "plan":
                        f.write(f"\n• 计划时间: {entry['timestamp']}\n")
                        f.write("  计划内容:\n")
                        for i, op in enumerate(entry["plan"], 1):
                            f.write(f"    {i}. {op['source']} -> {op['target']}\n")
                    elif entry["type"] == "execute":
                        f.write(f"\n• 执行时间: {entry['timestamp']}\n")
                        f.write(f"  执行结果: 成功 {entry['success_count']}/{entry['total']} 操作\n")
                        for result in entry["results"]:
                            f.write(f"    • {result.replace('✅ ', '').replace('⚠️ ', '').replace('❌ ', '')}\n")
            
            return f"📄 报告已生成: {report_filename}\n💡 路径: {report_path}"
        except Exception as e:
            return f"❌ 生成报告时出错: {str(e)}"

class CustomLLM(ChatOpenAI):
    """自定义LLM类，适配不同的大模型API"""
    
    def __init__(self, api_key: str, api_base: str, model_name: str, **kwargs):
        # 手动构建参数字典，避免传递不兼容的参数
        init_params = {
            "api_key": api_key,
            "base_url": api_base,
            "model": model_name
        }
        
        # 只添加已知兼容的参数
        compatible_params = ['temperature', 'max_tokens', 'timeout', 'max_retries']
        for param in compatible_params:
            if param in kwargs:
                init_params[param] = kwargs[param]
        
        super().__init__(**init_params)

def select_working_directory() -> str:
    """让用户选择工作目录"""
    home_dir = os.path.expanduser("~")
    common_dirs = [
        home_dir,
        os.path.join(home_dir, "Documents"),
        os.path.join(home_dir, "Downloads"),
        os.path.join(home_dir, "Desktop"),
        os.path.join(home_dir, "Pictures"),
        os.path.join(home_dir, "Videos")
    ]
    
    print("\n" + "=" * 50)
    print("📁 请选择工作目录:")
    print("=" * 50)
    for i, directory in enumerate(common_dirs, 1):
        print(f"{i}. {directory}")
    print(f"{len(common_dirs)+1}. 自定义路径")
    
    while True:
        try:
            choice = input(f"\n请输入选择 (1-{len(common_dirs)+1}): ")
            if not choice:
                continue
                
            choice = int(choice)
            if 1 <= choice <= len(common_dirs):
                return common_dirs[choice-1]
            elif choice == len(common_dirs)+1:
                while True:
                    custom_path = input("请输入自定义路径: ").strip()
                    if not custom_path:
                        continue
                        
                    if os.path.isdir(custom_path):
                        return custom_path
                    else:
                        print(f"❌ 错误: 路径 '{custom_path}' 不存在或不是目录")
                        retry = input("是否重新输入? (y/n): ")
                        if retry.lower() not in ['y', 'yes', '是']:
                            return select_working_directory()  # 递归返回重新选择
            else:
                print(f"❌ 无效选择，请输入 1-{len(common_dirs)+1} 的数字")
        except ValueError:
            print("❌ 请输入数字")

def setup_api_connection() -> tuple:
    """设置API连接信息"""
    print("\n" + "=" * 50)
    print("🔌 配置大语言模型API")
    print("=" * 50)
    
    # 从环境变量获取配置
    env_api_base = os.getenv('API_BASE')
    env_model_name = os.getenv('MODEL_NAME')
    env_api_key = os.getenv('API_KEY')
    
    if env_api_base and env_model_name and env_api_key:
        print("ℹ️  已从 .env 文件加载API配置")
        print(f"   API_BASE: {env_api_base}")
        print(f"   MODEL_NAME: {env_model_name}")
        confirm = input("是否使用这些配置? (y/n, 默认为y): ").strip().lower()
        if confirm in ['', 'y', 'yes', '是']:
            return env_api_key, env_api_base, env_model_name
    
    api_base = input(f"请输入API基础URL (例如 https://api.example.com/v1) [{env_api_base or '无默认值'}]: ").strip()
    if not api_base:
        api_base = env_api_base if env_api_base else "https://api.example.com/v1"  # 默认值
    
    model_name = input(f"请输入模型名称 (例如 qwen-max, gpt-4, mistral-large) [{env_model_name or '无默认值'}]: ").strip()
    if not model_name:
        model_name = env_model_name if env_model_name else "qwen-max"  # 默认值
    
    api_key = getpass(f"请输入API密钥 [{ '已设置' if env_api_key else '未设置'}]: ").strip()
    if not api_key:
        if env_api_key:
            api_key = env_api_key
        else:
            print("❌ 错误: API密钥不能为空")
            return setup_api_connection()
    
    return api_key, api_base, model_name

def setup_langchain_agent(api_key: str, api_base: str, model_name: str, base_dir: str) -> AgentExecutor:
    """设置LangChain代理"""
    try:
        # 创建自定义LLM实例
        llm = CustomLLM(
            api_key=api_key,
            api_base=api_base,
            model_name=model_name,
            temperature=0.3
        )
        
        mcp = FileSystemMCP(base_dir)
        
        # 创建工具
        tools = [
            Tool(
                name="list_directory",
                func=lambda path="": mcp.list_directory(path),
                description="列出指定路径的目录内容。参数: path (可选，默认为当前目录)"
            ),
            Tool(
                name="get_file_structure",
                func=lambda path="": mcp.get_file_structure(path=path, max_depth=3),
                description="获取目录结构（递归）。参数: path (可选路径)"
            ),
            Tool(
                name="plan_organize",
                func=mcp.plan_organize,
                description="制定文件整理计划。参数: organization_plan (JSON格式的整理计划)"
            ),
            Tool(
                name="execute_organize",
                func=mcp.execute_organize,
                description="执行已制定的整理计划。参数: confirm (确认执行, yes/no)"
            ),
            # 使用lambda包装generate_report以避免参数问题
            Tool(
                name="generate_report",
                func=lambda: mcp.generate_report(),
                description="生成整理报告。无需参数"
            )
        ]
        
        # 创建提示模板
        system_message = f"""你是一个专业的文件整理助手，帮助用户整理和组织他们的文件系统。

你的工作流程必须严格遵守:
1. 首先使用get_file_structure工具了解当前文件结构
2. 与用户充分沟通，明确整理需求和规则
3. 提出一个详细的整理计划，使用plan_organize工具记录计划
4. 详细解释计划内容，等待用户明确确认
5. 如果用户确认，使用execute_organize工具执行整理
6. 最后，使用generate_report工具生成整理报告

重要安全规则:
- ✅ 绝对不能删除任何文件，只能移动或重命名
- ✅ 在执行任何操作前，必须获得用户明确确认
- ✅ 所有整理计划必须是JSON格式，包含source和target字段的列表
- ✅ 必须使用实际存在的文件名，不能编造文件名
- ✅ target路径可以包含文件夹，会自动创建

当前工作目录: {base_dir}

输出格式要求:
- 使用清晰的结构化输出
- 重要信息使用✅/⚠️/❌等符号标记
- 提供操作前后的对比
- 执行前必须显示"等待确认"状态
"""
        
        prompt = ChatPromptTemplate.from_messages([
            ("system", system_message),
            MessagesPlaceholder(variable_name="chat_history"),
            ("user", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad")
        ])
        
        # 创建代理
        agent = create_tool_calling_agent(llm, tools, prompt)
        agent_executor = AgentExecutor(
            agent=agent, 
            tools=tools, 
            verbose=True,
            handle_parsing_errors=True,
            max_iterations=15  # 防止无限循环
        )
        
        return agent_executor, mcp
        
    except Exception as e:
        print(f"❌ 初始化LangChain代理失败: {str(e)}")
        print("💡 提示: 请检查API连接信息是否正确，以及网络连接是否正常")
        sys.exit(1)

def main():
    """主程序入口"""
    print("=" * 60)
    print("🚀 文件整理助手 - 基于大语言模型的智能文件组织系统")
    print("💡 通过自然语言指令，安全整理您的文件系统")
    print("=" * 60)
    
    # 步骤1: 选择工作目录
    base_dir = select_working_directory()
    print(f"\n📁 已选择工作目录: {base_dir}")
    
    # 检查目录是否可读
    if not os.access(base_dir, os.R_OK):
        print(f"❌ 错误: 没有权限读取目录 '{base_dir}'")
        print("💡 请确保您有权限访问该目录，或选择其他目录")
        return
    
    # 步骤2: 配置API连接
    api_key, api_base, model_name = setup_api_connection()
    
    # 步骤3: 设置LangChain代理
    print("\n⚙️ 正在初始化AI助手...")
    agent_executor, mcp = setup_langchain_agent(api_key, api_base, model_name, base_dir)
    print("✅ AI助手已准备就绪!")
    
    # 步骤4: 交互式对话
    print("\n" + "=" * 60)
    print("现在您可以与AI助手对话，描述您的文件整理需求。")
    print("示例提示: '请帮我整理这个目录，将图片文件按年份分类，文档文件按项目分类'")
    print("💡 提示: AI会先了解当前文件结构，然后提出整理计划，需要您确认后才会执行")
    print("=" * 60 + "\n")
    
    chat_history = []
    first_message = True
    
    while True:
        if first_message:
            user_input = "你好，我想整理这个目录下的文件，请帮我。"
            first_message = False
        else:
            user_input = input("\n您: ").strip()
            if not user_input:
                continue
                
            if user_input.lower() in ["exit", "quit", "退出", "bye"]:
                print("\n感谢使用文件整理助手! 再见!")
                break
        
        try:
            print("\n🧠 AI助手正在思考...")
            
            # 执行代理
            response = agent_executor.invoke({
                "input": user_input,
                "chat_history": chat_history
            })
            
            print("\n🤖 AI助手: " + response["output"])
            print("-" * 60)
            
            # 更新聊天历史
            chat_history.append(("human", user_input))
            chat_history.append(("ai", response["output"]))
            
            # 检查是否完成了整理并需要生成报告（更精确匹配成功消息）
            if "✅ 整理操作已完成!" in response["output"]:
                print("\n📊 正在生成整理报告...")
                report_result = mcp.generate_report()
                print("📝 " + report_result)
                print("-" * 60)
                
        except Exception as e:
            print(f"\n❌ 发生错误: {str(e)}")
            import traceback
            traceback.print_exc()
            print("💡 请重试或检查您的输入。")
            print("-" * 60)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n程序已由用户中断。")
        print("感谢使用文件整理助手!")
    except Exception as e:
        print(f"\n❌ 程序发生严重错误: {str(e)}")
        print("💡 建议检查:")
        print("   1. LangChain及相关依赖是否已正确安装")
        print("   2. API连接信息是否正确")
        print("   3. 工作目录权限是否正常")
        sys.exit(1)